import os

import numpy as np
import torch
from torchline.data import DATASET_REGISTRY

__all__ = [
    'MyDataset',
    'LSTMDataset'
]


@DATASET_REGISTRY.register()
class LSTMDataset(torch.utils.data.Dataset):
    def __init__(self, cfg):
        super(LSTMDataset, self).__init__()
        self.cfg = cfg
        dataset_cfg = cfg.dataset
        is_train = self.cfg.dataset.is_train
        data_path = os.path.join(dataset_cfg.dir, 'up_2dim.npy')
        label_path = os.path.join(dataset_cfg.dir, 'down_2dim.npy')
        data = np.load(data_path) # num*2*512
        labels = np.load(label_path)
        num_data = data.shape[0]
        # split train and test set
        data_size = int(num_data * 0.8)
        if is_train:
            self.data = data[:data_size, :, :]
            self.labels = labels[:data_size, :, :]
        else:
            self.data = data[data_size:, :, :]
            self.labels = labels[data_size:, :, :]
        real_size = len(self.labels)
        self.data = torch.from_numpy(self.data).float().view(real_size, 1, -1) # num*1*1024
        self.labels = torch.from_numpy(self.labels).float().view(real_size, -1) # num*1024

    def __getitem__(self, index):
        return self.data[index], self.labels[index]

    def __len__(self):
        return len(self.data)


@DATASET_REGISTRY.register()
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, cfg):
        super(MyDataset, self).__init__()
        self.cfg = cfg
        dataset_cfg = cfg.dataset
        is_train = self.cfg.dataset.is_train
        data_path = os.path.join(dataset_cfg.dir, 'up_2dim.npy')
        label_path = os.path.join(dataset_cfg.dir, 'down_2dim.npy')
        data = np.load(data_path) # num*2*512
        labels = np.load(label_path)
        num_data = data.shape[0]
        # split train and test set
        data_size = int(num_data * 0.8)
        if is_train:
            self.data = data[:data_size, :, :]
            self.labels = labels[:data_size, :, :]
        else:
            self.data = data[data_size:, :, :]
            self.labels = labels[data_size:, :, :]
        real_size = len(self.labels)
        self.data = torch.from_numpy(self.data).float() # num*2*512
        self.labels = torch.from_numpy(self.labels).float().view(real_size, -1) # num*1024

    def __getitem__(self, index):
        return self.data[index], self.labels[index]

    def __len__(self):
        return len(self.data)

if __name__ == "__main__":
    from torchline.config import get_cfg
    cfg = get_cfg()
    cfg.merge_from_file('./config/config.yaml')
    dataset = MyDataset(cfg)
    print(dataset)